from dna.Parameters import *
from dna.PolymorphismTables import PolymorphismTables
import dna.SiteFrequencySpectrum as sfs
import numpy.matlib as np
from math import sqrt
import scipy.stats.mstats as mstats
import dna.infiniteSitesTools as ist
import matplotlib.pyplot as plt
from dna.PolyspectOperator import *
from dna.PowerTools import Power
from easygui import *
import sys

msg = "Enter parameters below"
title = "Polyspect optimizer for CC model"
fieldNames = ["N(0) - Present effective size", \
			"N(t) - New population size", \
			"t - Time of pop. change (coalescent time)", \
			"S - Segregating sites", \
			"n - Number of Chromosomes", \
			"Number of simulations"]
defaultValues = ["10000", \
			"1000", \
			"0.25", \
			"50", \
			"6", \
			"1000"]  # we start with blanks for the values
fieldValues = multenterbox(msg,title, fieldNames, defaultValues)

# make sure that none of the fields was left blank
while 1:
	if fieldValues == None: break
	errmsg = ""
	for i in range(len(fieldNames)):
		if fieldValues[i].strip() == "":
			errmsg = errmsg + ('"%s" is a required field.\n\n' % fieldNames[i])
	if errmsg == "": break # no problems found
	fieldValues = multenterbox(errmsg, title, fieldNames, fieldValues)

#test
N = int( fieldValues[0].strip() )
NEW_SIZE = int( fieldValues[1].strip() )
t = float( fieldValues[2].strip() )
S = int( fieldValues[3].strip() )
NUM_OF_CHROMOSOMES = int( fieldValues[4].strip() )
SIM_SIZE = int( fieldValues[5].strip() )

RATIO = pow(10, 0.2)
SIM_SIZE = 1000
POWER_SIZE = SIM_SIZE / 10

LENGTH = 25000
MU = 1e-8

p = Population(NUM_OF_CHROMOSOMES, N, LENGTH, MU, 0.0, S)
d = Demography(N)
simParams = SimulationParameters( p, d )
print simParams

pt = PolymorphismTables()
msCommand = simParams.getMSString( SIM_SIZE )
msOut = pt.runMS(msCommand)
pt.readMS(None, msOut)

mySFSs = sfs.SiteFrequencySpectrum(pt)


msg = "Enter initial Polyspect below"
title = "Polyspect optimizer for CC model"
fieldNames = [str(i) for i in range(1,NUM_OF_CHROMOSOMES) ]
defaultValues = [str(1) for i in range(1,NUM_OF_CHROMOSOMES) ]

fieldValues = multenterbox(msg,title, fieldNames, defaultValues)
#
## make sure that none of the fields was left blank
#while 1:
#	if fieldValues == None: break
#	errmsg = ""
#	for i in range(len(fieldNames)):
#		if fieldValues[i].strip() == "":
#			errmsg = errmsg + ('"%s" is a required field.\n\n' % fieldNames[i])
#	if errmsg == "": break # no problems found
#	fieldValues = multenterbox(errmsg, title, fieldNames, fieldValues)


	

myTD = ist.getOp_TajimasD( NUM_OF_CHROMOSOMES )
TDpolyOp = Polyspect(myTD, simParams, SIM_SIZE, mySFSs )
ssList = TDpolyOp.operateOnAll( TDpolyOp.mySFSs )
print TDpolyOp.getRejectionValues(0.05)
myTDPower = Power(TDpolyOp, POWER_SIZE )

myFuLiD = ist.getOp_FuLiD( NUM_OF_CHROMOSOMES )
FLDpolyOp = Polyspect(myFuLiD, simParams, SIM_SIZE, mySFSs )
ssList = FLDpolyOp.operateOnAll( FLDpolyOp.mySFSs )
print FLDpolyOp.getRejectionValues(0.05)
myFLDPower = Power(FLDpolyOp, POWER_SIZE )

myFuLiF = ist.getOp_FuLiF( NUM_OF_CHROMOSOMES )
FLFpolyOp = Polyspect(myFuLiF, simParams, SIM_SIZE, mySFSs )
ssList = FLFpolyOp.operateOnAll( FLFpolyOp.mySFSs )
print FLFpolyOp.getRejectionValues(0.05)
myFLFPower = Power(FLFpolyOp, POWER_SIZE )

altDemography = Demography(N)
ep = altDemography.addEpoch(1000, NEW_SIZE, 0.0)
ep.setIterationParameters('time', 0.001/RATIO, RATIO, 1.0, True)

myTDPowers = []
myFLDPowers = []
myFLFPowers = []
myTimes = []
for e in ep:
	pt = PolymorphismTables()
	msCommand = simParams.getAlternativeMSString( POWER_SIZE, e )
	msOut = pt.runMS(msCommand)
	pt.readMS(None, msOut)
	myAltSFSs = sfs.SiteFrequencySpectrum(pt)
	
	TDpowerIs = myTDPower.getPower(altDemography, myAltSFSs)
	myTDPowers.append(TDpowerIs)
	
	FLDpowerIs = myFLDPower.getPower(altDemography, myAltSFSs)
	myFLDPowers.append(FLDpowerIs)
	
	FLFpowerIs = myFLFPower.getPower(altDemography, myAltSFSs)
	myFLFPowers.append(FLFpowerIs)
	
	myTimes.append(e.myTime)
	print TDpowerIs, FLDpowerIs, FLFpowerIs

for i in range(len(myTimes)):
	print '%1.4f \t%1.4f' % (myTimes[i], myTDPowers[i])

print '\n'

for i in range(len(myTimes)):
	print '%1.4f \t%1.4f' % (myTimes[i], myFLDPowers[i])

print '\n'

for i in range(len(myTimes)):
	print '%1.4f \t%1.4f' % (myTimes[i], myFLFPowers[i])

print '\n'

plt.semilogx( myTimes, myTDPowers, myTimes, myFLDPowers, myTimes, myFLFPowers )
plt.ylabel('Power')
plt.xlabel('Time of event (in coalescent time)')
plt.title('CC model with 100-fold growth. S=50, n=6')
plt.legend( ('Tajimas D', 'FuLi D', 'FuLi F') )
plt.ylim( 0.0, 1.0 )
plt.show()